pacman::p_load(tidyverse,haven,ggplot2,gridExtra,estimatr,ggeffects,cregg,expss,stargazer,arm,here,ggtext,texreg,contrast)

# colors c("#CB2314","#273046","#354823","#FAD510")

onecolour <- c("#004488")

twocolours <- c("#004488","#DDAA33")

threecolours <- c("#004488","#DDAA33","#BB5566")

fourcolours <- c("#004488","#DDAA33","#BB5566","#000000")

fivecolours <- c("#004488","#DDAA33","#BB5566","#000000","#228833")

###################
## Set up functions
###################

# Figure 1

plot_canvas_figure1 <- function(df = df.wide, q = q,y = ..prop..,labs = "label", ylimits = c(0, .45)){
  ggplot(data=df) + 
    geom_bar(mapping = aes(x = q, y = ..prop..), 
             stat = "count", width = .9,
             color=onecolour, fill = onecolour,
             size=1) +
    xlab("") + ylab("")+
    labs(title=labs) +
    scale_x_continuous(guide = guide_axis(n.dodge = 1), 
                       labels = function(x) stringr::str_wrap(c("Much too low", 
                                                                "Too low", 
                                                                "About right",
                                                                "Too high", 
                                                                "Much too high", 
                                                                "Don't know"), width = 6),
                       breaks = c(1,2,3,4,5,6)) +
    scale_y_continuous(limits = ylimits,labels = scales::percent_format(accuracy = 1)) +
    theme_minimal() +
    theme(axis.text=element_text(size=20),
          axis.title=element_text(size=20),
          title=element_text(size=22),
          panel.border = element_blank(), panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(), axis.line = element_line(colour = "black"))
}

# Coefficient plots

pd <- position_dodge(0.75)

# More models in one

plot_canvas_coefficient <- function(df = coefficients_logit, face = face, x = reorder(model, desc(model)), models = NA, cols = threecolours){
  ggplot(data = df, aes(y = estimate, x = x, color = models)) + 
    geom_point(position = pd,size = 2) + 
    geom_linerange(aes(ymin = estimate-1.96*stderror, ymax = estimate+1.96*stderror),
                   lwd = 1/2,
                   position = pd) +
    geom_linerange(aes(ymin = estimate-1.645*stderror, ymax = estimate+1.645*stderror),
                   lwd = 1,
                   position = pd) +
    coord_flip() +
    labs(x = "", 
         y = "Coefficient estimate",
         color = "") +
    geom_hline(yintercept = 0,
               linetype = "dashed") +
    theme_bw(base_size = 30) +
    theme(panel.border = element_blank(), panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(), axis.line = element_line(colour = "black"),
          legend.position = "bottom",
          axis.text.y = element_text(face = df$bold),   
          axis.title=element_text(size=25)) +   #coefcombined / tibblefigure
    scale_color_manual(values=cols,
                       guide = guide_legend(reverse = TRUE))# +      
  #  guides(col=guide_legend(ncol=1,nrow=length(unique(models))))
}

# Get coefficients

getcoef <- function(model_name = "i am a model", reg = reg){
  tibble(group = c(rep(model_name,14)),
         model=c("Value of own house \n (ref: Don't own):","£100k and under","£100k to £200k","£200k to £300k","£300k to £400k","£400k to £500k","£500k to £750k","£750k and up",
                 "Value of parents' house \n (ref: Don't own):","Not applicable/deceased","£200k and under","£200k to £400k","£400k to £600k","£600k and up"),
         estimate = c(NA,reg$estimate[2],reg$estimate[3],reg$estimate[4],reg$estimate[5],reg$estimate[6],reg$estimate[7],reg$estimate[8],
                      NA,reg$estimate[9],reg$estimate[10],reg$estimate[11],reg$estimate[12],reg$estimate[13]),
         stderror = c(NA,reg$std.error[2],reg$std.error[3],reg$std.error[4],reg$std.error[5],reg$std.error[6],reg$std.error[7],reg$std.error[8],
                      NA,reg$std.error[9],reg$std.error[10],reg$std.error[11],reg$std.error[12],reg$std.error[13]),
         bold = c(rep("plain",5),"bold",rep("plain",7),"bold")) %>% 
    mutate(model = factor(model, levels = model))
}

### Vignette experiment

getcoef_experiment_conditional <- function(model_name = "i am a model", reg = reg, facettext = "Overall level"){
  tibble(group = model_name,
         model=c("Death/Double tax","Equality of opportunity","Taxes and public goods"),
         estimate = c(reg$Contrast[1],reg$Contrast[2],reg$Contrast[3]),
         stderror = c(reg$SE[1],reg$SE[2],reg$SE[3]),
         bold = c(rep("plain",3)),
         facettext = facettext) %>% 
    mutate(model = factor(model, levels = model))
}

contrast_conditional <- function(reg = ownreg,facettext = "Title",parents = FALSE, party_ID = FALSE){
  
  if(parents == FALSE & party_ID == FALSE){
  
  own1_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
    reg = contrast(reg,
                   list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), homeowner = "Don't Own", info_treatment = "Control"),
                   list(framing_treatment=c("Control"), homeowner =  "Don't Own", info_treatment = "Control"),
                   type = "individual"),
    model_name = "Don't own",
    facettext = facettext)
  
  own2_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
    reg = contrast(reg,
                   list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), homeowner = "Own", info_treatment = "Control"),
                   list(framing_treatment=c("Control"), homeowner =  "Own", info_treatment = "Control"),
                   type = "individual"),
    model_name = "Own",
    facettext = facettext)
  
  cond <- rbind(own1_dk,own2_dk)
  
  }
  
  if(parents == TRUE & party_ID == FALSE){
    
    own1_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), house_rent_groups_parents2 = "Don't Own", info_treatment = "Control"),
                     list(framing_treatment=c("Control"), house_rent_groups_parents2 = "Don't Own", info_treatment = "Control"),
                     type = "individual"),
      model_name = "Don't own",
      facettext = facettext)
    
    own2_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), house_rent_groups_parents2 = "Own", info_treatment = "Control"),
                     list(framing_treatment=c("Control"), house_rent_groups_parents2 = "Own", info_treatment = "Control"),
                     type = "individual"),
      model_name = "Own",
      facettext = facettext)
  
    own3_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), house_rent_groups_parents2 = "Not applicable/deceased", info_treatment = "Control"),
                     list(framing_treatment=c("Control"), house_rent_groups_parents2 = "Not applicable/deceased", info_treatment = "Control"),
                     type = "individual"),
      model_name = "Not applicable/deceased",
      facettext = facettext)
    
  cond <- rbind(own1_dk,own2_dk,own3_dk)
    
  }
  
  if(parents == FALSE & party_ID == TRUE){
    
    own1_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), party_ID = "Conservative", info_treatment = "Control"),
                     list(framing_treatment=c("Control"), party_ID = "Conservative", info_treatment = "Control"),
                     type = "individual"),
      model_name = "Conservative",
      facettext = facettext)
    
    own2_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), party_ID = "Labour", info_treatment = "Control"),
                     list(framing_treatment=c("Control"), party_ID = "Labour", info_treatment = "Control"),
                     type = "individual"),
      model_name = "Labour",
      facettext = facettext)
    
    own3_dk <- getcoef_experiment_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(framing_treatment=c("Double Taxation","Levelling Playing Field","Spending and Tax Cuts"), party_ID = "Other/Don't know", info_treatment = "Control"),
                     list(framing_treatment=c("Control"), party_ID = "Other/Don't know", info_treatment = "Control"),
                     type = "individual"),
      model_name = "Other/Don't know",
      facettext = facettext)
    
    cond <- rbind(own1_dk,own2_dk,own3_dk)
    
  }
  
  return(cond)
  
}

# Information experiment

getcoef_information_conditional <- function(model_name = "i am a model", reg = reg,facettext = "Overall level"){
  tibble(group = model_name,
         model=c("National","National+Local"),
         estimate = c(reg$Contrast[1],reg$Contrast[2]),
         stderror = c(reg$SE[1],reg$SE[2]),
         bold = c(rep("plain",2)),
         facettext = facettext) %>% 
    mutate(model = factor(model, levels = model))
}

contrast_information_conditional <- function(reg = ownreg,facettext = "Title",parents = FALSE){
  
  if(parents == FALSE){
    
    own1_dk <- getcoef_information_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(Qtreatment=c("2","3"), house_rent_2groups = "Don't Own"),
                     list(Qtreatment=c("1"), house_rent_2groups = "Don't Own"),
                     type = "individual"),
      model_name = "Don't own",
      facettext = facettext)
    
    own2_dk <- getcoef_information_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(Qtreatment=c("2","3"), house_rent_2groups = "Own"),
                     list(Qtreatment=c("1"), house_rent_2groups = "Own"),
                     type = "individual"),
      model_name = "Own",
      facettext = facettext)
    
    cond <- rbind(own1_dk,own2_dk)
    
  }
  
  if(parents == TRUE){
    
    own1_dk <- getcoef_information_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(Qtreatment=c("2","3"), house_rent_3groups_parents = "Don't Own"),
                     list(Qtreatment=c("1"), house_rent_3groups_parents = "Don't Own"),
                     type = "individual"),
      model_name = "Don't own",
      facettext = facettext)
    
    own2_dk <- getcoef_information_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(Qtreatment=c("2","3"), house_rent_3groups_parents = "Own"),
                     list(Qtreatment=c("1"), house_rent_3groups_parents = "Own"),
                     type = "individual"),
      model_name = "Own",
      facettext = facettext)
    
    own3_dk <- getcoef_information_conditional( # Helper-function to get the coefficients
      reg = contrast(reg,
                     list(Qtreatment=c("2","3"), house_rent_3groups_parents = "Not applicable/deceased"),
                     list(Qtreatment=c("1"), house_rent_3groups_parents = "Not applicable/deceased"),
                     type = "individual"),
      model_name = "Not applicable/deceased",
      facettext = facettext)
    
    cond <- rbind(own1_dk,own2_dk,own3_dk)
    
  }
  
  return(cond)
  
}

# Get coefficient experiment

getcoef_experiment <- function(model_name = "i am a model", reg = reg){
  tibble(group = model_name,
         model=c("Death/Double tax","Equality of opportunity","Taxes and public goods"),
         estimate = c(reg$coefficients[2],reg$coefficients[3],reg$coefficients[4]),
         stderror = c(reg$std.error[2],reg$std.error[3],reg$std.error[4]),
         bold = c(rep("plain",3))) %>% 
    mutate(model = factor(model, levels = model))
}


# Plot (both experiments)

# for parents to work x = factor(group2, level = c("Own", "NA/\ndeceased", "Don't own"))
plot_canvas_coefficient_experiment <-  function(df = coefficients_logit, givemebreaks = c("Taxes and public goods","Equality of opportunity","Death/Double tax")){
  ggplot(data = df, aes(y = estimate, x = group, color = model)) + 
    geom_point(position = pd,size = 2) + 
    geom_linerange(aes(ymin = estimate-1.96*stderror, ymax = estimate+1.96*stderror),
                   lwd = 1/2,
                   position = pd) +
    geom_linerange(aes(ymin = estimate-1.645*stderror, ymax = estimate+1.645*stderror),
                   lwd = 1,
                   position = pd) +
    coord_flip() +
    labs(x = "", 
         y = "Coefficient estimate",
         color = "") +
    geom_hline(yintercept = 0,
               linetype = "dashed") +
    theme_bw(base_size = 30) +
    theme(panel.border = element_blank(), panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(), axis.line = element_line(colour = "black"),
          legend.position = "bottom",
          legend.key.size = unit(1.5, 'cm'),
          legend.text=element_text(size=25),
          text = element_text(size = 30),
          title=element_text(size=35),
          axis.title=element_text(size=25)) +
    scale_color_manual(values= threecolours,
                       guide = guide_legend(reverse = TRUE),
                       breaks = givemebreaks) +
    facet_wrap(~ facettext , nrow =1) 
}

plot_canvas_coefficient_experiment2 <- function(df = coefficients_logit, face = face, x = reorder(model, desc(model)), models = NA, cols = c("#273046","#354823","#FAD510","Black")){
  ggplot(data = df, aes(y = estimate, x = x, color = models)) + 
    geom_point(position = pd,size = 2) + 
    geom_linerange(aes(ymin = estimate-1.96*stderror, ymax = estimate+1.96*stderror),
                   lwd = 1/2,
                   position = pd) +
    geom_linerange(aes(ymin = estimate-1.645*stderror, ymax = estimate+1.645*stderror),
                   lwd = 1,
                   position = pd) +
    coord_flip() +
    labs(x = "", 
         y = "Coefficient estimate",
         color = "") +
    geom_hline(yintercept = 0,
               linetype = "dashed") +
    theme_bw(base_size = 23) +
    theme(panel.border = element_blank(), panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(), axis.line = element_line(colour = "black"),
          legend.position = "right",
          axis.title=element_text(size=18)) +   #coefcombined
    scale_color_manual(values=cols,
                       guide = guide_legend(reverse = TRUE))# +      
  #  guides(col=guide_legend(ncol=1,nrow=length(unique(models))))
}